package algorithms;

import java.awt.Color;
import java.awt.Component;
import java.awt.Graphics;
import java.util.Stack;

public class HuffmanTree {

	// Tree root pointer
	private Node root;
	private Node nullNode;

	// Tree node
	public class Node {
		private Node left;
		private Node right;
		private Node ancestor;
		private Character character;
		private int amount = 0;
		private int height = 1;

		// Node constructors
		public Node() {
			character = null;
		}

		public Node(Character newCharacter) {
			character = newCharacter;
		}

		public Character getCharacter() {
			return character;
		}
	}

	// Tree constructor
	public HuffmanTree() {
		nullNode = new Node();
		nullNode.amount = 0;
		// nullNode.character = (char)-1;
	}

	// Fixes heights
	private void fixHeight() {
		Stack<Node> stack = new Stack<Node>();
		stack.push(root);
		while (!stack.empty()) {
			Node node = stack.pop();
			if (node.left != null) {
				node.left.height = node.height + 1;
				// node.left.ancestor = node;
				stack.push(node.left);
			}
			if (node.right != null) {
				node.right.height = node.height + 1;
				// node.right.ancestor = node;
				stack.push(node.right);
			}
		}
	}

	// Increases amounts
	private void incAncestors(Node node) {
		while (node != null) {
			node.amount++;
			node = node.ancestor;
		}
	}

	// Decreases amounts
	private void decAncestors(Node node) {
		while (node != null) {
			node.amount--;
			node = node.ancestor;
		}
	}

	// Searches for the character and returns it's node
	private Node find(Character c) {
		if (root != null) {
			Stack<Node> s = new Stack<Node>();
			Node t = root;
			s.push(t);
			if (root.character == c)
				return root;
			else {
				while (!s.empty()) {
					t = s.pop();
					if (t.character == c)
						return t;
					if (t.left != null)
						s.push(t.left);
					if (t.right != null)
						s.push(t.right);
				}
			}
		}
		return null;
	}

	// Swaps nodes
	private void swap(Node node1, Node node2) {
		// Store ancestor of node1
		Node node = node1.ancestor;

		// Change pointer to node1
		if (node1.ancestor.left == node1)
			node1.ancestor.left = node2;
		else
			node1.ancestor.right = node2;

		// Change ancestor of node1
		node1.ancestor = node2.ancestor;

		// Change pointer to node2
		if (node2.ancestor.left == node2)
			node2.ancestor.left = node1;
		else
			node2.ancestor.right = node1;

		// Change ancestor of node2
		node2.ancestor = node;

		// Swap heights
		int height = node1.height;
		node1.height = node2.height;
		node2.height = height;
		fixHeight();
	}

	// Finds min amount
	private Node findMinOnHeght(int height) {
		Node min = new Node();
		min.amount = 9100500;
		Stack<Node> stack = new Stack<Node>();
		if (root == null)
			return null;
		stack.push(root);
		while (!stack.empty()) {
			Node node = stack.pop();
			if (node.height == height)
				if (node.amount < min.amount)
					min = node;
			if (node.left != null)
				stack.push(node.left);
			if (node.right != null)
				stack.push(node.right);
		}
		return min;
	}

	// Finds max amount
	private Node findMaxOnHeght(int height) {
		Node max = new Node();
		max.amount = 0;
		Stack<Node> stack = new Stack<Node>();
		if (root == null)
			return null;
		stack.push(root);
		while (!stack.empty()) {
			Node node = stack.pop();
			if (node.height == height)
				if (node.amount > max.amount)
					max = node;
			if (node.left != null)
				stack.push(node.left);
			if (node.right != null)
				stack.push(node.right);
		}
		return max;
	}

	// Balances the tree
	private void balance() {
		for (int i = 2; i < nullNode.height - 1; i++) {
			Node min = findMinOnHeght(i);
			Node max = findMaxOnHeght(i + 1);
			if (min.amount < max.amount) {
				swap(min, max);
				incAncestors(max.ancestor);
				decAncestors(min.ancestor);
				i = 1;
			}
		}
	}

	// Node to draw
	private class PaintNode {
		private Node node;
		private int x, y, size = 22;

		public PaintNode(Node newNode, int x, int y) {
			node = newNode;
			this.x = x;
			this.y = y;
		}

		private void drawNode(Graphics g) {
			if (node == null)
				return;
			if (node.height < 7) {
				if (node.character != null) {
					g.drawRect(x - size / 2, y - size / 2, size, size);
					if (node.character == (char) 18)
						g.drawString("\'0\':" + node.amount, x - size / 2 + 2,
								y + 5);
					else if (node.character == (char) 16)
						g.drawString("\'1\':" + node.amount, x - size / 2 + 2,
								y + 5);
					else
						g.drawString("\'" + node.character + "\':"
								+ node.amount, x - size / 2 + 2, y + 5);
				} else {
					if (node == nullNode) {
						g.drawRect(x - size / 2 - 1, y - size / 2 - 1,
								size + 1, size + 1);
						g.setColor(Color.white);
						g.fillRect(x - size / 2, y - size / 2, size, size);
						g.setColor(Color.black);
					} else {
						g.drawOval(x - size / 2, y - size / 2, size, size);
						g.drawString(node.amount + "", x - 5, y + 5);
					}
				}
			} else {
				if (node.character != null) {
					g.drawLine(x, y, x, y + 20);
				}
				size = 1;
			}

		}
	}

	private void printTree(Graphics g, PaintNode node, int shift) {
		if (node.node == null)
			return;
		node.drawNode(g);
		double Y = node.y + node.size + 40;
		Node leftNode = node.node.left;
		Node rightNode = node.node.right;
		if (leftNode != null) {
			double leftX = node.x - shift;
			PaintNode pN = new PaintNode(leftNode, (int) leftX, (int) Y);
			if (pN.node.height == 7) {
				g.drawLine(node.x - node.size / 2, node.y, pN.x, pN.y);
			} else if (pN.node.height > 7) {
				g.drawLine(node.x, node.y, pN.x, pN.y);
				node.size = 1;
			} else
				g.drawLine(node.x - node.size / 2, node.y, pN.x, pN.y - pN.size
						/ 2);
			printTree(g, pN, shift / 2);
		}
		if (rightNode != null) {
			double rightX = node.x + shift;
			PaintNode pN = new PaintNode(rightNode, (int) rightX, (int) Y);
			if (pN.node.height == 7) {
				g.drawLine(node.x + node.size / 2, node.y, pN.x, pN.y);
			} else if (pN.node.height > 7) {
				g.drawLine(node.x, node.y, pN.x, pN.y);
				node.size = 1;
			} else
				g.drawLine(node.x + node.size / 2, node.y, pN.x, pN.y - pN.size
						/ 2);
			printTree(g, pN, shift / 2);
		}
	}

	public Node getNullNode() {
		return nullNode;
	}

	// Adds element
	public void add(Character character) {
		Node node = find(character);
		if (node == null) {
			node = nullNode.ancestor;
			Node ancestor = new Node();
			if (node == null) {
				node = new Node();
				root = ancestor;
			}
			ancestor.left = nullNode;
			ancestor.right = new Node(character);
			ancestor.right.ancestor = ancestor;
			if (ancestor != root)
				ancestor.ancestor = node;
			node.left = ancestor;
			nullNode.ancestor = ancestor;
			incAncestors(ancestor.right);
		} else
			incAncestors(node);
		balance();
		fixHeight();
	}

	// Returns character code
	public String getCharCode(Character c) {
		if (c == '0')
			c = (char) 18;
		if (c == '1')
			c = (char) 16;
		if (nullNode.ancestor == null)
			return "" + c;
		String string = new String();
		Node node = find(c);
		if (node == null) {
			node = nullNode;
			string += c;
		}
		while (node != root) {
			if (node.ancestor.right == node)
				string = '1' + string;
			else
				string = '0' + string;
			node = node.ancestor;
		}
		return string;
	}

	// returns null if code is not full or returns Node with character or
	// returns nullNode if code is full but it's first time appearance of
	// character
	public Node getCharByCode(String code) {
		Node tempNode = root;
		for (int i = 0; i < code.length(); i++) {
			if (code.charAt(i) == '0')
				tempNode = tempNode.left;
			else
				tempNode = tempNode.right;
		}
		if (tempNode.character == null)
			return null;
		return tempNode;
	}

	// Prints the tree
	public void printTree(Component component, int offset) {
		Graphics g = component.getGraphics();
		printTree(g, new PaintNode(root, component.getWidth() / 2, offset), 250);
	}

}